'''
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.


Copyright (c) 2009 bryan p. conrad
bconrad2@gmail.com

Created on: 7/26/2009

Version: 0.1

@author: bryan
'''
#===============================================================================
# Imports
#===============================================================================
import JT 

# for keyboard interaction
from msvcrt import kbhit, getch
# To stop render window form non-responding
import win32gui     


class DRR:
    """
    A collection of components that can be used to create a digitally
    reconstructed radiograph (DRR).

    With a decent graphics card the rendering can be hardware accelerated and does
    not necessarily need to be done on-screen.  A WindowToImageFilter is used 
    to capture the contents of the render window and save the projection as a 2D image.
    """

    def __init__(self):
        """
        Initialize DRR Class.
        """
        #Set up Moving Image 3D renderwindow, with interaction
        self._ren = JT.vtk.vtkRenderer()
        self._renWin = JT.vtk.vtkRenderWindow()
        self._renWin.AddRenderer(self._ren)
        self._iren = JT.vtk.vtkRenderWindowInteractor()
        self._iren.SetRenderWindow(self._renWin)
        self._ren.SetLayer(0)
        self.SetBlackOnWhiteImage()
        
        # Set default render window size 
        # To be set in the registration class once the size of the fixed 
        # image is known
        #self._renWin.SetSize(512,512)
        self._xrayCamera = JT.XRayCamera()
        self._ren.SetActiveCamera(self._xrayCamera)

        #We will need to use a custom interaction scheme (based on JointTrack)
        style = JT.JTInteractorStyle(self._iren)
        self._iren.SetInteractorStyle(style)

        #Create filter to capture an image from the render window
        self._w2if = JT.vtk.vtkWindowToImageFilter()
        self._w2if.SetInput(self._renWin)

        #Create a VTK Image Flipper
        self._vtkFlipperX = JT.vtk.vtkImageFlip()
        self._vtkFlipperY = JT.vtk.vtkImageFlip()

        self._vtkFlipperX.SetFilteredAxis(0)
        self._vtkFlipperY.SetFilteredAxis(1)

        #Create a VTK cast image filter to convert to float and char
        self._floatCast = JT.vtk.vtkImageCast()
        self._uCharCast = JT.vtk.vtkImageCast()

        self._floatCast.SetOutputScalarTypeToFloat()
        self._floatCast.ClampOverflowOn()

        self._uCharCast.SetOutputScalarTypeToUnsignedChar()
        self._uCharCast.ClampOverflowOn()

        # Fixed image size
        self._imageSize = []

        # Result of Generating an image is a projected image (2D VTK image)
        self._projectedImage = []

        # Create components for a background image
        self._renWin.SetNumberOfLayers(2)

        self._imageMapper = JT.vtk.vtkImageMapper()
        self._imageActor = JT.vtk.vtkActor2D()
        self._imageRen = JT.vtk.vtkRenderer()

        self._renWin.AddRenderer(self._imageRen)
        self._imageActor.SetMapper(self._imageMapper)
        self._imageRen.AddActor2D(self._imageActor)
        self._imageRen.SetLayer(1)
        self._imageRen.SetInteractive(0)

        self._backgroundImage = None

        self._pose = (0,0,0,0,0,0)
        self._volume_list = []
        self._volume = None

        
    def AddVolume(self,volume):
        """
        DEPRECATED.  Use SetVolume instead.  A list of volumes is not maintained
        only one volume can be loaded at a time.  All volumes are managed
        by the Data class, not by the DRR.  
        """
        print 'DRR.AddVolume is Deprecated.  Use SetVolume instead.'
        self.SetVolume(volume)
        
    def RemoveAllVolumes(self):
        """
        Remove all JT volumes (and any other props) from the renderer.
        """
        self._ren.RemoveAllViewProps()
        self._volume_list = []
        
    def SetVolume(self,volume):
        """
        Set a JT volume in the renderer.
        """
        self.RemoveAllVolumes()
        self._volume = volume
        self._ren.AddVolume(volume)
    
    def AddBodyCenter(self,center):
        self._volume.AddBodyCenter(center)
        
    def AddSphereAtBodyCenter(self,center):
        # Add axes to renderer
        sphere = JT.vtk.vtkSphereSource()
        sphere.SetRadius(5.0)
        sphere.SetPhiResolution(10)
        mapper = JT.vtk.vtkPolyDataMapper()
        sphere_actor = JT.vtk.vtkActor()
        mapper.SetInput(sphere.GetOutput())
        sphere_actor.SetMapper(mapper)
        sphere_actor.SetPosition(center)
        props = sphere_actor.GetProperty()
        props.SetColor(1,0,0)
        self._ren.AddActor(sphere_actor)
    
    def SetRenderWindowSize(self,size):
        """
        size is a tuple of integers that represents the size of the render
        window in pixels. 
        """
        self._renWin.SetSize(size)
    
    def Render(self):
        """
        Display renderwindow.  Does not start interaction.

        TODO: need to check to make sure that everything is setup.
        """
        self._renWin.Render()
        currentWindowName = self._renWin.GetWindowName()
        if currentWindowName[0:4] == 'Vis':
            self._renWin.SetWindowName("JointTrack - DRR")
        
    def StartInteraction(self):
        """
        Start the render window interactor.  If used in a script, the user must 
        quit the interaction (press q, or e typically) to continue to the 
        next command.
        """
        self._iren.Start()
    
    def UseVTKInteractionStyle(self):
        """
        Warning: This style of interaction will change the calibration setup,
        By allowing the user to move the position of the camera.  Do not use this 
        function for image registration.
        
        This method will turn on the VTK interactor style trackball camera.
        """
        style = JT.vtk.vtkInteractorStyleTrackballCamera()
        self._iren.SetInteractorStyle(style)
        
        
    def GenerateImage(self):
        """
        Render the volume and capture the resulting image from the render window.
        A change in this version of jt is to NOT convert from VTK image to ITK image
        """
        #To avoid being recognized as a non-responding program in windows, 
        #we need to check in on the renderwindow at each iteration
        win32gui.PeekMessage(0,0,0,0)
        
        #Capture projection from render window
        self._renWin.Render()
        self._w2if.Modified()

        #for image to be displayed correctly, it has to be flipped about the X axis
        self._vtkFlipperX.SetInput(self._w2if.GetOutput())
        self._vtkFlipperY.SetInput(self._vtkFlipperX.GetOutput())

        # CHANGE - projected image is stored as vtk format (no ITK)
        self._projectedImage = self._vtkFlipperY.GetOutput()
        self._projectedImage.Update()
        
        self._pyProjectedImage = JT.utils.vtk2numpy(self._projectedImage)

    def SaveGeneratedImage(self,filename):
        """
        Save the generated image to disk.  If a projection does not exist,
        it will be generated and saved, otherwise this function will save the
        image currently residing in self._projectedImage.

        The appropriate vtk image file writer is used, as determined 
        by the filename extension.
        """
        #print "Generating Image..."
        self.GenerateImage()
        JT.utils.genericImageWriter(self._projectedImage, filename)
        
                        
    def SetBlackOnWhiteImage(self):
        """
        This function changes the backgound color of the renderer.
        The renderer seems to automatically change the color of the actor so that
        it shows up on the background.
        """
        self._ren.SetBackground(1,1,1)
        self._blackOnWhiteFlag = 1

    def SetWhiteOnBlackImage(self):
        """
        This function changes the backgound color of the renderer.
        The renderer seems to automatically change the color of the actor so that
        it shows up on the background.

        This is the default behavior.
        """
        self._ren.SetBackground(0,0,0)
        self._blackOnWhiteFlag = 0

    def GetXRayCamera(self):
        """
        Return the active Camera for the scene
        """
        return self._xrayCamera

    def SetXRayCamera(self,camera):
        """
        Set the active Camera for the scene
        """
        self._xrayCamera = camera
        self._ren.SetActiveCamera(camera)
        
        #Determine whether an internal or external calibration setup is being used
        try:
            cal=self._xrayCamera._IntCal
            cal.isValid()
        except:
            try:
                cal=self._xrayCamera._ExtCal
                cal.isValid()
            except:
                logging.debug("DRR.GenerateImage: Cannot Generate Image, No calibration set")
                logging.debug("External Cal: %s, Internal Cal: %s" %
                               (self._xrayCamera._IntCal.isValid(),self._xrayCamera._ExtCal.isValid()))

    def GetInternalCalibration(self):
        """
        The internal Calibration currently being used for the active camera
        """
        return self._xrayCamera._IntCal

    def GetExternalCalibration(self):
        """
        The external Calibration currently being used for the active camera
        """
        return self._xrayCamera._ExtCal
    
    def SetExternalCalibration(self,cal):
        self._xrayCamera.SetExternalCalibration(cal)

    def GetImage(self):
        """
        Return the current projected image (VTK format),
        The python (numpy) version can be obtained by using self._pyProjectedImage
        """
        #Assume that the user wants a new image to be generated.
        self.GenerateImage()
        return self._projectedImage
        
    def GetRenderTime(self):
        """
        Get the last render time in seconds (from renderer)
        """
        return self._ren.GetLastRenderTimeInSeconds()

    def SetBackgroundImage(self,vtkImage):
        """
        Supply a VKT image to be displayed as the background in the same render
        window with the DRR projection.
        """
        #if a background image is used, imageRen needs to be moved to the back layer
        self._ren.SetLayer(1)
        self._imageRen.SetLayer(0)

        self._backgroundImage = vtkImage

        self._imageMapper.SetInput(self._backgroundImage)

        (minPixel, maxPixel) = self._backgroundImage.GetScalarRange()

        self._imageMapper.SetColorLevel(maxPixel/2)
        self._imageMapper.SetColorWindow(maxPixel)

        self.Render()

    def InteractionOff(self):
        """
        Turn interaction off, useful for optimizing registration.  After user
        has specified a starting position
        """
        self._renWin.SetInteractor(None)

    def InteractionOn(self):
        """
        Turn interaction on. For user to specify a starting position
        """
        self._iren.SetRenderWindow(self._renWin)
        